import pandas as pd
import numpy as np

# vote_rating.py 是生成同名.csv文件的程序，生成了投票数和得分相关的表格

summer_movies = pd.read_csv("DSProject/csv/summer_movies.csv")

vote_rating = summer_movies[["tconst", "average_rating", "num_votes"]].copy()
vote_rating = vote_rating.sort_values("average_rating")
# 划分投票数区间（大于30000的电影归为一个区间上）
votes_bins = (
    list(1000 * np.arange(10))
    + list(10000 * np.arange(1, 3))
    + [int(vote_rating["num_votes"].max() + 1)]
)
# 划分更细致的投票数区间（聚焦在低投票数上）
votes_bins_detailed = list(100 * np.arange(11)) + [
    int(vote_rating["num_votes"].max() + 1)
]

# 划分得分区间
rating_bins = np.arange(11)
# 划分更细致的得分区间（便于后面对高分电影分析）
rating_bins_detailed = np.linspace(0, 10, 21)
# 进行列表列添加
vote_rating["votes_group"] = pd.cut(vote_rating["num_votes"], votes_bins)
vote_rating["rating_group"] = pd.cut(vote_rating["average_rating"], rating_bins)
vote_rating["rating_group_detailed"] = pd.cut(
    vote_rating["average_rating"], rating_bins_detailed
)
vote_rating["votes_group_detailed"] = pd.cut(
    vote_rating["num_votes"], votes_bins_detailed
)
vote_rating.to_csv("DSProject/csv/vote_rating.csv", index=False)
